/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

#pragma once

#include <string>
#include <unordered_map>
#include <unordered_set>

#include <json/json.h>

#include <mariana-trench/Assert.h>
#include <mariana-trench/Context.h>
#include <mariana-trench/IncludeMacros.h>
#include <mariana-trench/Kind.h>
#include <mariana-trench/Rule.h>
#include <mariana-trench/TransformList.h>

namespace marianatrench {

/**
 * Represents a rule source -> sink rule that is additionally constrained on
 * a call-chain flow from the root callable of the source -> sink flow to any
 * call-chain effect source.
 * e.g.
 *                    exported_callable(): with effect_source: Exported
 *                            |
 *                    root_callable(): with inferred effect_sink: Source@Sink
 *                            |
 *          +-----------------+-----------------+
 *          |                                   |
 * source_callable(): with              sink_callable(): with
 *  source kind: Source                    sink kind: Sink
 *
 * Effective rule is then internally modelled as:
 *  Exported -> Source -> Sink
 */
class SourceSinkWithExploitabilityRule final : public Rule {
 public:
  using KindToTransformsMap =
      std::unordered_map<const Kind*, const TransformList*>;

 public:
  SourceSinkWithExploitabilityRule(
      const std::string& name,
      int code,
      const std::string& description,
      KindSet effect_source_kinds,
      KindSet source_kinds,
      KindSet sink_kinds,
      KindToTransformsMap source_as_transforms)
      : Rule(name, code, description),
        effect_source_kinds_(std::move(effect_source_kinds)),
        source_kinds_(std::move(source_kinds)),
        sink_kinds_(std::move(sink_kinds)),
        source_as_transforms_(std::move(source_as_transforms)) {}

  DELETE_COPY_CONSTRUCTORS_AND_ASSIGNMENTS(SourceSinkWithExploitabilityRule)

  const KindSet& effect_source_kinds() const {
    return effect_source_kinds_;
  }

  const KindSet& source_kinds() const {
    return source_kinds_;
  }

  const KindSet& sink_kinds() const {
    return sink_kinds_;
  }

  bool uses(const Kind*) const override;

  const TransformList* MT_NULLABLE
  source_as_transform(const Kind* source_kind) const;

  std::optional<CoveredRule> coverage(
      const KindSet& sources,
      const KindSet& sinks,
      const TransformSet& transforms) const override;

  static std::unique_ptr<Rule> from_json(
      const std::string& name,
      int code,
      const std::string& description,
      const Json::Value& value,
      Context& context);
  Json::Value to_json() const override;

 private:
  KindSet effect_source_kinds_;
  KindSet source_kinds_;
  KindSet sink_kinds_;
  KindToTransformsMap source_as_transforms_;
};

} // namespace marianatrench
